package com.ai.aichat.controller;

import com.ai.aichat.adopter.AliAiServiceAdopter;
import com.ai.aichat.adopter.IAIServiceAdopter;
import com.ai.aichat.adopter.OpenAiServiceAdopter;
import com.ai.aichat.dto.ChatImage;
import com.ai.aichat.dto.ChatMessage;
import com.fasterxml.jackson.core.JsonProcessingException;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * @author lwlianghehe@gmail.com
 * @date 2025/9/15
 */
@RestController
@RequestMapping("/api/ai")
public class AiController {
    private final AliAiServiceAdopter aliAiServiceAdopter;
    private final OpenAiServiceAdopter openAiServiceAdopter;

    private final Map<String, IAIServiceAdopter> aiServiceAdopterMap = new HashMap<>();

    private IAIServiceAdopter getAIServiceAdopter(String platform) {
        if (platform.equals("ali")) {
            return aliAiServiceAdopter;
        }

        return openAiServiceAdopter;
    }

    public AiController(AliAiServiceAdopter aliAiServiceAdopter, OpenAiServiceAdopter openAiServiceAdopter) {
        this.aliAiServiceAdopter = aliAiServiceAdopter;
        this.openAiServiceAdopter = openAiServiceAdopter;
        aiServiceAdopterMap.put("ali", aliAiServiceAdopter);
        aiServiceAdopterMap.put("openai", openAiServiceAdopter);
    }

    @PostMapping("/chat/text")
    public ChatMessage sendText(@RequestBody Map<String, Object> params) throws JsonProcessingException {
        String model = (String) params.get("model");
        String message = (String) params.get("message");
        String platform = (String) params.get("platform");

//        return getAIServiceAdopter(platform).sendText(model, message);
        return aiServiceAdopterMap.get(platform).sendText(model, message);
    }

    @PostMapping("/chat/text/stream")
    public SseEmitter sendTextStream(@RequestBody Map<String, Object> params) throws JsonProcessingException {
        String model = (String) params.get("model");
        String message = (String) params.get("message");
        String platform = (String) params.get("platform");

        // 0 表示不超时；生产可根据需要设置超时并做心跳保活
        SseEmitter emitter = new SseEmitter(0L);
        // 将阻塞式上游读取放入异步，避免占用 Servlet 线程
        CompletableFuture.runAsync(() -> {
            try {
                aiServiceAdopterMap.get(platform).sendText(model, message, true, delta -> {
                    try {
                        emitter.send(SseEmitter.event().name("delta").data(delta));
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                });

                // 上游完成后通知前端结束
                emitter.send(SseEmitter.event().name("done").data("[DONE]"));
                emitter.complete();
            } catch (Exception ex) {
                emitter.completeWithError(ex);
            }
        });
        return emitter;
    }

    @PostMapping("create/image")
    public ChatImage createImage(@RequestBody Map<String, Object> params) throws JsonProcessingException {
        String model = (String) params.get("model");
        String prompt = (String) params.get("prompt");
        String size = (String) params.get("size");
        String platform = (String) params.get("platform");

//        return getAIServiceAdopter(platform).createImage(model, prompt, 1, size);
        return aiServiceAdopterMap.get(platform).createImage(model, prompt, 1, size);
    }
}
